#!/usr/bin/env python3
"""
模型下载脚本
自动下载AI动画生成系统所需的预训练模型
"""

import os
import sys
import requests
import zipfile
import tarfile
from pathlib import Path
from tqdm import tqdm
import logging
from huggingface_hub import hf_hub_download, snapshot_download
import torch

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

logger = logging.getLogger(__name__)


class ModelDownloader:
    """模型下载器"""
    
    def __init__(self, models_dir="models"):
        self.models_dir = Path(models_dir)
        self.models_dir.mkdir(parents=True, exist_ok=True)
        
        # 模型配置
        self.models = {
            "stable_diffusion": {
                "name": "runwayml/stable-diffusion-v1-5",
                "type": "diffusers",
                "description": "Stable Diffusion v1.5 文本到图像模型",
                "size": "4.2GB"
            },
            "stable_diffusion_xl": {
                "name": "stabilityai/stable-diffusion-xl-base-1.0",
                "type": "diffusers",
                "description": "Stable Diffusion XL 基础模型",
                "size": "6.9GB"
            },
            "text_to_video": {
                "name": "damo-vilab/text-to-video-ms-1.7b",
                "type": "diffusers",
                "description": "文本到视频生成模型",
                "size": "3.8GB"
            },
            "stylegan3": {
                "name": "NVlabs/stylegan3",
                "type": "git",
                "description": "StyleGAN3 风格生成模型",
                "size": "2.1GB"
            },
            "video_llama": {
                "name": "DAMO-NLP-SG/Video-LLaMA",
                "type": "transformers",
                "description": "Video-LLaMA 视频理解模型",
                "size": "1.5GB"
            },
            "clip": {
                "name": "openai/clip-vit-base-patch32",
                "type": "transformers",
                "description": "CLIP 视觉语言模型",
                "size": "0.5GB"
            },
            "bert": {
                "name": "bert-base-uncased",
                "type": "transformers",
                "description": "BERT 文本编码模型",
                "size": "0.4GB"
            },
            "bert_chinese": {
                "name": "bert-base-chinese",
                "type": "transformers",
                "description": "中文BERT模型",
                "size": "0.4GB"
            }
        }
    
    def download_file(self, url: str, filepath: Path, description: str = ""):
        """下载文件"""
        try:
            response = requests.get(url, stream=True)
            response.raise_for_status()
            
            total_size = int(response.headers.get('content-length', 0))
            
            with open(filepath, 'wb') as f:
                with tqdm(total=total_size, unit='B', unit_scale=True, 
                         desc=description) as pbar:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                            pbar.update(len(chunk))
            
            logger.info(f"下载完成: {filepath}")
            return True
            
        except Exception as e:
            logger.error(f"下载失败: {e}")
            return False
    
    def download_huggingface_model(self, model_name: str, model_dir: Path):
        """从Hugging Face下载模型"""
        try:
            logger.info(f"正在下载Hugging Face模型: {model_name}")
            
            # 使用snapshot_download下载整个模型
            snapshot_download(
                repo_id=model_name,
                local_dir=model_dir,
                local_dir_use_symlinks=False
            )
            
            logger.info(f"Hugging Face模型下载完成: {model_name}")
            return True
            
        except Exception as e:
            logger.error(f"Hugging Face模型下载失败: {e}")
            return False
    
    def download_diffusers_model(self, model_name: str, model_dir: Path):
        """下载Diffusers模型"""
        try:
            logger.info(f"正在下载Diffusers模型: {model_name}")
            
            # 创建模型目录
            model_dir.mkdir(parents=True, exist_ok=True)
            
            # 下载模型文件
            snapshot_download(
                repo_id=model_name,
                local_dir=model_dir,
                local_dir_use_symlinks=False,
                ignore_patterns=["*.md", "*.txt", "*.git*"]
            )
            
            logger.info(f"Diffusers模型下载完成: {model_name}")
            return True
            
        except Exception as e:
            logger.error(f"Diffusers模型下载失败: {e}")
            return False
    
    def download_transformers_model(self, model_name: str, model_dir: Path):
        """下载Transformers模型"""
        try:
            logger.info(f"正在下载Transformers模型: {model_name}")
            
            # 创建模型目录
            model_dir.mkdir(parents=True, exist_ok=True)
            
            # 下载模型文件
            snapshot_download(
                repo_id=model_name,
                local_dir=model_dir,
                local_dir_use_symlinks=False,
                ignore_patterns=["*.md", "*.txt", "*.git*"]
            )
            
            logger.info(f"Transformers模型下载完成: {model_name}")
            return True
            
        except Exception as e:
            logger.error(f"Transformers模型下载失败: {e}")
            return False
    
    def download_stylegan3(self, model_dir: Path):
        """下载StyleGAN3模型"""
        try:
            logger.info("正在下载StyleGAN3模型")
            
            # StyleGAN3的下载链接（示例）
            stylegan3_urls = {
                "stylegan3-t-ffhqu-1024x1024.pkl": "https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhqu-1024x1024.pkl",
                "stylegan3-r-ffhqu-1024x1024.pkl": "https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhqu-1024x1024.pkl"
            }
            
            model_dir.mkdir(parents=True, exist_ok=True)
            
            for filename, url in stylegan3_urls.items():
                filepath = model_dir / filename
                if not filepath.exists():
                    success = self.download_file(url, filepath, f"StyleGAN3 {filename}")
                    if not success:
                        logger.warning(f"StyleGAN3 {filename} 下载失败")
                else:
                    logger.info(f"StyleGAN3 {filename} 已存在")
            
            logger.info("StyleGAN3模型下载完成")
            return True
            
        except Exception as e:
            logger.error(f"StyleGAN3模型下载失败: {e}")
            return False
    
    def download_model(self, model_key: str, force_download: bool = False):
        """下载指定模型"""
        if model_key not in self.models:
            logger.error(f"未知模型: {model_key}")
            return False
        
        model_config = self.models[model_key]
        model_name = model_config["name"]
        model_type = model_config["type"]
        model_dir = self.models_dir / model_key
        
        # 检查是否已存在
        if model_dir.exists() and not force_download:
            logger.info(f"模型已存在: {model_key}")
            return True
        
        logger.info(f"开始下载模型: {model_key} ({model_config['description']})")
        logger.info(f"模型大小: {model_config['size']}")
        
        success = False
        
        if model_type == "diffusers":
            success = self.download_diffusers_model(model_name, model_dir)
        elif model_type == "transformers":
            success = self.download_transformers_model(model_name, model_dir)
        elif model_type == "git":
            if model_key == "stylegan3":
                success = self.download_stylegan3(model_dir)
            else:
                success = self.download_huggingface_model(model_name, model_dir)
        else:
            success = self.download_huggingface_model(model_name, model_dir)
        
        if success:
            logger.info(f"模型下载成功: {model_key}")
        else:
            logger.error(f"模型下载失败: {model_key}")
        
        return success
    
    def download_all_models(self, force_download: bool = False):
        """下载所有模型"""
        logger.info("开始下载所有模型...")
        
        results = {}
        total_models = len(self.models)
        
        for i, model_key in enumerate(self.models.keys(), 1):
            logger.info(f"进度: {i}/{total_models} - {model_key}")
            
            try:
                success = self.download_model(model_key, force_download)
                results[model_key] = success
            except Exception as e:
                logger.error(f"下载模型 {model_key} 时出错: {e}")
                results[model_key] = False
        
        # 统计结果
        successful = sum(results.values())
        failed = len(results) - successful
        
        logger.info(f"下载完成: 成功 {successful} 个，失败 {failed} 个")
        
        if failed > 0:
            logger.warning("失败的模型:")
            for model_key, success in results.items():
                if not success:
                    logger.warning(f"  - {model_key}")
        
        return results
    
    def check_model_status(self):
        """检查模型状态"""
        logger.info("检查模型状态...")
        
        status = {}
        for model_key, model_config in self.models.items():
            model_dir = self.models_dir / model_key
            
            if model_dir.exists():
                # 检查文件大小
                total_size = sum(f.stat().st_size for f in model_dir.rglob('*') if f.is_file())
                size_mb = total_size / (1024 * 1024)
                
                status[model_key] = {
                    "exists": True,
                    "size_mb": round(size_mb, 2),
                    "description": model_config["description"]
                }
            else:
                status[model_key] = {
                    "exists": False,
                    "size_mb": 0,
                    "description": model_config["description"]
                }
        
        return status
    
    def print_status(self):
        """打印模型状态"""
        status = self.check_model_status()
        
        print("\n模型状态:")
        print("-" * 80)
        print(f"{'模型名称':<20} {'状态':<10} {'大小(MB)':<12} {'描述'}")
        print("-" * 80)
        
        for model_key, info in status.items():
            status_text = "✓ 已下载" if info["exists"] else "✗ 未下载"
            size_text = f"{info['size_mb']:.1f}" if info["exists"] else "0.0"
            
            print(f"{model_key:<20} {status_text:<10} {size_text:<12} {info['description']}")
        
        print("-" * 80)
        
        # 统计信息
        total_models = len(status)
        downloaded_models = sum(1 for info in status.values() if info["exists"])
        total_size = sum(info["size_mb"] for info in status.values())
        
        print(f"总计: {downloaded_models}/{total_models} 个模型已下载，总大小: {total_size:.1f} MB")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="AI动画生成系统模型下载器")
    parser.add_argument("--model", type=str, help="下载指定模型")
    parser.add_argument("--all", action="store_true", help="下载所有模型")
    parser.add_argument("--status", action="store_true", help="检查模型状态")
    parser.add_argument("--force", action="store_true", help="强制重新下载")
    parser.add_argument("--models-dir", type=str, default="models", help="模型存储目录")
    
    args = parser.parse_args()
    
    # 设置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    
    # 创建下载器
    downloader = ModelDownloader(args.models_dir)
    
    if args.status:
        # 检查状态
        downloader.print_status()
    elif args.model:
        # 下载指定模型
        success = downloader.download_model(args.model, args.force)
        if success:
            print(f"✓ 模型 {args.model} 下载成功")
        else:
            print(f"✗ 模型 {args.model} 下载失败")
    elif args.all:
        # 下载所有模型
        results = downloader.download_all_models(args.force)
        successful = sum(results.values())
        total = len(results)
        print(f"\n下载完成: {successful}/{total} 个模型成功")
    else:
        # 显示帮助
        parser.print_help()
        print("\n可用模型:")
        for model_key, model_config in downloader.models.items():
            print(f"  {model_key}: {model_config['description']} ({model_config['size']})")


if __name__ == "__main__":
    main() 